import triton
import triton.language as tl

from torch._inductor.runtime import triton_helpers
from torch._inductor.runtime.triton_helpers import libdevice

import torch
import torch_npu
from .util import triton_helper_compile, compiled_kernel_list

# triton_per_fused_add_embedding_mean_mul_pow_rsqrt_10
@triton.jit
def fused_kernel11(in_out_ptr0, in_ptr0, in_ptr1, in_ptr2, in_ptr3, in_ptr4, out_ptr1, xnumel, rnumel):
    xnumel = 1
    XBLOCK: tl.constexpr = 1
    rnumel = 2048
    RBLOCK: tl.constexpr = 2048
    xoffset = tl.program_id(0) * XBLOCK
    xindex = tl.full([1], xoffset, tl.int32)
    xmask = tl.full([RBLOCK], True, tl.int1)
    rindex = tl.arange(0, RBLOCK)[:]
    roffset = 0
    rmask = tl.full([RBLOCK], True, tl.int1)
    r0 = rindex
    tmp0 = tl.load(in_ptr0 + (0))
    tmp1 = tl.broadcast_to(tmp0, [RBLOCK])
    tmp8 = tl.load(in_out_ptr0 + (r0), None)
    tmp10 = tl.load(in_ptr2 + (r0), None)
    tmp12 = tl.load(in_ptr3 + (r0), None)
    tmp18 = tl.load(in_ptr4 + (r0), None)
    tmp2 = tl.full([RBLOCK], 128256, tl.int32)
    tmp3 = tmp1 + tmp2
    tmp4 = tmp1 < 0
    tmp5 = tl.where(tmp4, tmp3, tmp1)
    tl.device_assert((0 <= tmp5) & (tmp5 < 128256), "index out of bounds: 0 <= tmp5 < 128256")
    tmp7 = tl.load(in_ptr1 + (r0 + (2048*tmp5)), None)
    tmp9 = tmp7 + tmp8
    tmp11 = tmp9 + tmp10
    tmp13 = tmp11 + tmp12
    tmp14 = tmp13 * tmp13
    tmp15 = tl.broadcast_to(tmp14, [RBLOCK])
    tmp17 = triton_helpers.promote_to_tensor(tl.sum(tmp15, 0))
    tmp19 = 2048.0
    tmp20 = tmp17 / tmp19
    tmp21 = 1e-05
    tmp22 = tmp20 + tmp21
    tmp23 = libdevice.rsqrt(tmp22)
    tmp24 = tmp13 * tmp23
    tmp25 = tmp18 * tmp24
    tl.store(in_out_ptr0 + (tl.broadcast_to(r0, [RBLOCK])), tmp13, None)
    tl.store(out_ptr1 + (tl.broadcast_to(r0, [RBLOCK])), tmp25, None)

def test_compile_for_fused_kernel11():
    signature = {0: '*fp32', 1: '*i32', 2: '*fp32', 3: '*fp32', 4: '*fp32', 5: '*fp32', 6: '*fp32', 7: 'i32', 8: 'i32'}
    triton_helper_compile(func=fused_kernel11, signature=signature)

compiled_kernel_list["fused_kernel11"] = test_compile_for_fused_kernel11